# coding:utf-8
# utils/network_fusion.py

import networkx as nx
import numpy as np
import os


def build_network_from_edge_file(file_path, layer_name="layer"):
    """
    从边文件中构建一个 NetworkX 图对象。
    :param file_path: 边文件路径，格式为 u v weight（每行三个字段）
    :param layer_name: 图的名称（用于标识不同网络层）
    :return: nx.Graph 对象
    """
    G = nx.Graph(name=layer_name)
    with open(file_path, 'r') as f:
        for line in f:
            if line.startswith('#'):
                continue
            parts = line.strip().split()
            if len(parts) >= 3:
                u, v, w = parts[0], parts[1], float(parts[2])
                G.add_edge(u, v, weight=w)
    return G


def build_multilayer_networks(ppi_graph, expr_graph, mapping_dict=None):
    """
    构建多层图并添加跨层映射边。
    :param ppi_graph: PPI 层图 (nx.Graph)
    :param expr_graph: 基因表达层图 (nx.Graph)
    :param mapping_dict: 字典 {protein_id: gene_id} 表示蛋白质与基因之间的映射
    :return: 多层图（包含 interlayer=True 标记的层间边）
    """
    # 初始化综合图
    fused_G = nx.Graph()

    # 添加PPI层内的边
    for u, v, data in ppi_graph.edges(data=True):
        fused_G.add_edge(f"ppi:{u}", f"ppi:{v}", weight=data["weight"], layer="PPI")

    # 添加Expression层内的边
    for u, v, data in expr_graph.edges(data=True):
        fused_G.add_edge(f"expr:{u}", f"expr:{v}", weight=data["weight"], layer="Expression")

    # 如果存在层间映射，则添加跨层边
    if mapping_dict:
        for protein_id, gene_id in mapping_dict.items():
            if str(protein_id) in ppi_graph.nodes() and str(gene_id) in expr_graph.nodes():
                fused_G.add_edge(f"ppi:{protein_id}", f"expr:{gene_id}", weight=0.5, interlayer=True)

    return fused_G


def fuse_two_layers(ppi_edges, expr_edges, alpha=0.5, output_file="data/processed/fused_network.txt"):
    """
    将两个网络层融合为一个综合网络。
    :param ppi_edges: PPI 网络边列表 [(u, v, weight), ...]
    :param expr_edges: 基因表达网络边列表 [(u, v, weight), ...]
    :param alpha: 权重系数，控制 PPI 和 GE 的贡献比例
    :param output_file: 输出融合后的边文件路径
    :return: 融合后的 NetworkX 图对象
    """
    combined_graph = nx.Graph()

    # 合并 PPI 和 Expression 网络的节点集合
    all_nodes = set()
    for u, v, w in ppi_edges + expr_edges:
        all_nodes.add(u)
        all_nodes.add(v)

    for node in all_nodes:
        combined_graph.add_node(node)

    # 加入 PPI 边
    for u, v, w in ppi_edges:
        combined_graph.add_edge(u, v, weight=alpha * float(w))

    # 加入 Gene Expression 边
    for u, v, w in expr_edges:
        combined_graph.add_edge(u, v, weight=(1 - alpha) * float(w))

    # 可选：合并重复边（如果同一个节点对出现在多个网络中）
    merged_edges = {}
    for u, v, data in combined_graph.edges(data=True):
        key = tuple(sorted((u, v)))
        merged_edges[key] = data['weight']

    # 清空并重新加入合并后的边
    combined_graph.clear_edges()
    for (u, v), weight in merged_edges.items():
        combined_graph.add_edge(u, v, weight=weight)

    # 保存到文件
    save_graph_to_file(combined_graph, output_file)

    return combined_graph


def compute_combined_weight(ppi_weight, expr_weight, alpha=0.5):
    """
    计算综合边权重：
    wcombined(u,v) = α * w_ppi(u,v) + (1-α) * w_ge(u,v)
    """
    return alpha * ppi_weight + (1 - alpha) * expr_weight


def save_graph_to_file(G, file_path):
    """
    将 NetworkX 图保存为边列表文本文件。
    :param G: nx.Graph 对象
    :param file_path: 输出文件路径
    """
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    with open(file_path, "w") as f:
        for u, v, data in G.edges(data=True):
            f.write(f"{u}\t{v}\t{data['weight']:.4f}\n")


def load_unified_network(file_path):
    """
    读取统一格式的 Entrez ID 网络。
    :param file_path: 统一边界文件路径（如 unified_ppi.txt）
    :return: nx.Graph
    """
    G = nx.Graph()
    with open(file_path, 'r') as f:
        for line in f:
            if line.startswith('#'): continue
            u, v, w = line.strip().split()
            G.add_edge(int(u), int(v), weight=float(w))
    return G


def main():
    """
    主函数：加载原始网络、融合并输出结果。
    """
    # 加载预处理后的 Entrez ID 数据
    ppi_net = load_unified_network("data/processed/unified_ppi.txt")
    expr_net = load_unified_network("data/processed/unified_gene_expression.txt")

    # 构建融合网络
    print("开始融合网络...")
    fused_network = nx.Graph()

    # 设置融合参数
    alpha = 0.6  # PPI 占比更高

    # 添加 PPI 网络的边（加权）
    for u, v, data in ppi_net.edges(data=True):
        fused_network.add_edge(u, v, weight=compute_combined_weight(data['weight'], 0, alpha))

    # 添加 Gene Expression 网络的边（加权）
    for u, v, data in expr_net.edges(data=True):
        if fused_network.has_edge(u, v):
            # 如果已有边，则更新为平均权重
            current_weight = fused_network[u][v]['weight']
            new_weight = current_weight + compute_combined_weight(0, data['weight'], alpha)
            fused_network[u][v]['weight'] = new_weight
        else:
            fused_network.add_edge(u, v, weight=compute_combined_weight(0, data['weight'], alpha))

    # 保存融合后的网络
    save_graph_to_file(fused_network, "results/fused_network.txt")
    print("网络融合完成，已保存至 results/fused_network.txt")


if __name__ == "__main__":
    main()